import random

import torch
from torchvision.datasets import MNIST
from torch.utils.data import Subset
from torchvision import transforms as T

# 这种数据处理的方式常用于
# 1.快速测试模型性能，加上你好数据量以加速训练
# 2.创建简化版数据集用于快速演示
# 3.当某种类别严格不能过多导致类别不平衡的时候

tf1 = T.Compose([
    T.Resize((32, 32), antialias=True),
    T.ToTensor()
])


def tf2(label):
    return torch.tensor(label)


# 1.加载数据集
ds = MNIST("xiazai",
           download=False,
           train=True,
           transform=tf1,
           target_transform=tf2
           )

# 2. 我现在只想训练识别数字 0 1 2， 其它不要，怎么做
# 初始化字典data，按类别0 1 2存储样本索引
# 根据存储的样本索引，利用Subset创建子集。

data = {
    0: [],
    1: [],
    2: []
}

# 3. 遍历数据集，收集label为0 1 2 的样本索引
for i in range(len(ds)):
    img, label = ds[i]
    if label.item() in [0, 1, 2]:
        data[label.item()].append(i)
for i in data:
    print(f"类别{i}的样本数量为：{len(data[i])}")
# 类别0的样本数量为：5923
# 类别1的样本数量为：6742
# 类别2的样本数量为：5958

# 4. 感觉5000张还是比较多，我还可以下采样 50%
rate = 0.5
for k, v in data.items():
    v_len = len(v)
    count = round(v_len * rate)
    data[k] = random.sample(v, count)

for i in data:
    print(f"类别{i}的样本数量为：{len(data[i])}")

# 5. 合并样本索引
idx = []
for v in data.values():
    idx += v

# 6. 截取子集
subset = Subset(ds, idx)

# 7. 保存
import pickle

"""
.pkl文件时python中用于序列化 和 反序列化对象的文本格式，通过import pickle实现·
1.pkl文件本质时二进制文件，可以存储python对象（列表、字典、模型、类等）状态信息
2.需要序列化就是存储  pickle.dump(data, f)   mode = ‘wb’
3.反序列化就是加载 data = pickle.load(d)   mode = 'rb'
"""
with open("subset.pkl", "wb") as f:
    pickle.dump(subset, f)


# pickle 序列化与反序列化的注意事项
# 1. 模块导入一致性
# 序列化和反序列化时必须保持相同的模块导入
# 保存数据时用到的类/函数在加载时必须可用且路径一致
# 如代码中的 tf2 函数必须在加载前定义
